#include "tee_client_api.h"
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

/* Simplifies calls to accept() */
/* $begin sockaddrdef */
typedef struct sockaddr SA;
/* $end sockaddrdef */

/* Persistent state for the robust I/O (Rio) package */
/* $begin rio_t */
#define RIO_BUFSIZE 8192
typedef struct {
  int rio_fd;                /* Descriptor for this internal buf */
  int rio_cnt;               /* Unread bytes in internal buf */
  char *rio_bufptr;          /* Next unread byte in internal buf */
  char rio_buf[RIO_BUFSIZE]; /* Internal buffer */
} rio_t;
/* $end rio_t */

/* External variables */
extern char **environ; /* Defined by libc */

/* Misc constants */
#define MAXLINE 8192 /* Max text line length */
#define LISTENQ 1024 /* Second argument to listen() */
#define MAX_READ_ONCE 1
#define OPERATION_START_FLAG 1

/* Signal wrappers */
typedef void handler_t(int);

/**************************
 * Error-handling functions
 **************************/

/* $begin errorfuns */
/* $begin unixerror */
void unix_error(char *msg) /* Unix-style error */
{
  fprintf(stderr, "%s: %s\n", msg, strerror(errno));
  exit(0);
}
/* $end unixerror */

/************************************
 * Wrappers for Unix signal functions
 ***********************************/

/* $begin sigaction */
handler_t *Signal(int signum, handler_t *handler) {
  struct sigaction action, old_action;

  action.sa_handler = handler;
  sigemptyset(&action.sa_mask); /* Block sigs of type being handled */
  action.sa_flags = SA_RESTART; /* Restart syscalls if possible */

  if (sigaction(signum, &action, &old_action) < 0)
    unix_error("Signal error");
  return (old_action.sa_handler);
}
/* $end sigaction */

/********************************
 * Wrappers for Unix I/O routines
 ********************************/

void Close(int fd) {
  int rc;

  if ((rc = close(fd)) < 0)
    unix_error("Close error");
}

/****************************
 * Sockets interface wrappers
 ****************************/

void Listen(int s, int backlog) {
  int rc;

  if ((rc = listen(s, backlog)) < 0)
    unix_error("Listen error");
}

int Accept(int s, struct sockaddr *addr, socklen_t *addrlen) {
  int rc;

  if ((rc = accept(s, addr, addrlen)) < 0)
    unix_error("Accept error");
  return rc;
}

/****************************************
 * The Rio package - Robust I/O functions
 ****************************************/

/*
 * rio_writen - Robustly write n bytes (unbuffered)
 */
/* $begin rio_writen */
ssize_t rio_writen(int fd, void *usrbuf, size_t n) {
  size_t nleft = n;
  ssize_t nwritten;
  char *bufp = usrbuf;

  while (nleft > 0) {
    if ((nwritten = write(fd, bufp, nleft)) <= 0) {
      if (errno == EINTR) /* Interrupted by sig handler return */
        nwritten = 0;     /* and call write() again */
      else
        return -1; /* errno set by write() */
    }
    nleft -= nwritten;
    bufp += nwritten;
  }
  return n;
}
/* $end rio_writen */

/*
 * rio_read - This is a wrapper for the Unix read() function that
 *    transfers min(n, rio_cnt) bytes from an internal buffer to a user
 *    buffer, where n is the number of bytes requested by the user and
 *    rio_cnt is the number of unread bytes in the internal buffer. On
 *    entry, rio_read() refills the internal buffer via a call to
 *    read() if the internal buffer is empty.
 */
/* $begin rio_read */
static ssize_t rio_read(rio_t *rp, char *usrbuf, size_t n) {
  int cnt;

  while (rp->rio_cnt <= 0) { /* Refill if buf is empty */
    rp->rio_cnt = read(rp->rio_fd, rp->rio_buf, sizeof(rp->rio_buf));
    if (rp->rio_cnt < 0) {
      if (errno != EINTR) /* Interrupted by sig handler return */
        return -1;
    } else if (rp->rio_cnt == 0) /* EOF */
      return 0;
    else
      rp->rio_bufptr = rp->rio_buf; /* Reset buffer ptr */
  }

  /* Copy min(n, rp->rio_cnt) bytes from internal buf to user buf */
  cnt = n;
  if (rp->rio_cnt < n)
    cnt = rp->rio_cnt;
  memcpy(usrbuf, rp->rio_bufptr, cnt);
  rp->rio_bufptr += cnt;
  rp->rio_cnt -= cnt;
  return cnt;
}
/* $end rio_read */

/*
 * rio_readinitb - Associate a descriptor with a read buffer and reset buffer
 */
/* $begin rio_readinitb */
void rio_readinitb(rio_t *rp, int fd) {
  rp->rio_fd = fd;
  rp->rio_cnt = 0;
  rp->rio_bufptr = rp->rio_buf;
}
/* $end rio_readinitb */

/*
 * rio_readnb - Robustly read n bytes (buffered)
 */
/* $begin rio_readnb */
ssize_t rio_readnb(rio_t *rp, void *usrbuf, size_t n) {
  size_t nleft = n;
  ssize_t nread;
  char *bufp = usrbuf;

  while (nleft > 0) {
    if ((nread = rio_read(rp, bufp, nleft)) < 0)
      return -1; /* errno set by read() */
    else if (nread == 0)
      break; /* EOF */
    nleft -= nread;
    bufp += nread;
  }
  return (n - nleft); /* return >= 0 */
}
/* $end rio_readnb */

/*
 * rio_readlineb - Robustly read a text line (buffered)
 */
/* $begin rio_readlineb */
ssize_t rio_readlineb(rio_t *rp, void *usrbuf, size_t maxlen) {
  int n, rc;
  char c, *bufp = usrbuf;

  for (n = 1; n < maxlen; n++) {
    if ((rc = rio_read(rp, &c, 1)) == 1) {
      *bufp++ = c;
      if (c == '\n') {
        n++;
        break;
      }
    } else if (rc == 0) {
      if (n == 1)
        return 0; /* EOF, no data read */
      else
        break; /* EOF, some data was read */
    } else
      return -1; /* Error */
  }
  *bufp = 0;
  return n - 1;
}
/* $end rio_readlineb */

/**********************************
 * Wrappers for robust I/O routines
 **********************************/

void Rio_readinitb(rio_t *rp, int fd) { rio_readinitb(rp, fd); }

/********************************
 * Client/server helper functions
 ********************************/
/*
 * open_clientfd - Open connection to server at <hostname, port> and
 *     return a socket descriptor ready for reading and writing. This
 *     function is reentrant and protocol-independent.
 *
 *     On error, returns:
 *       -2 for getaddrinfo error
 *       -1 with errno set for other errors.
 */
/* $begin open_clientfd */
int open_clientfd(char *hostname, char *port) {
  int clientfd, rc;
  struct addrinfo hints, *listp, *p;

  /* Get a list of potential server addresses */
  memset(&hints, 0, sizeof(struct addrinfo));
  hints.ai_socktype = SOCK_STREAM; /* Open a connection */
  hints.ai_flags = AI_NUMERICSERV; /* ... using a numeric port arg. */
  hints.ai_flags |= AI_ADDRCONFIG; /* Recommended for connections */
  if ((rc = getaddrinfo(hostname, port, &hints, &listp)) != 0) {
    fprintf(stderr, "getaddrinfo failed (%s:%s): %s\n", hostname, port,
            gai_strerror(rc));
    return -2;
  }

  /* Walk the list for one that we can successfully connect to */
  for (p = listp; p; p = p->ai_next) {
    /* Create a socket descriptor */
    if ((clientfd = socket(p->ai_family, p->ai_socktype, p->ai_protocol)) < 0) {
      int reuse = 1;
      setsockopt(clientfd, SOL_SOCKET, SO_REUSEADDR, (void *)&reuse,
                 sizeof(reuse));
      continue; /* Socket failed, try the next */
    }

    /* Connect to the server */
    if (connect(clientfd, p->ai_addr, p->ai_addrlen) != -1)
      break; /* Success */
    if (close(clientfd) <
        0) { /* Connect failed, try another */ // line:netp:openclientfd:closefd
      fprintf(stderr, "open_clientfd: close failed: %s\n", strerror(errno));
      return -1;
    }
  }

  /* Clean up */
  freeaddrinfo(listp);
  if (!p) /* All connects failed */
    return -1;
  else /* The last connect succeeded */
    return clientfd;
}
/* $end open_clientfd */

/*
 * open_listenfd - Open and return a listening socket on port. This
 *     function is reentrant and protocol-independent.
 *
 *     On error, returns:
 *       -2 for getaddrinfo error
 *       -1 with errno set for other errors.
 */
/* $begin open_listenfd */
int open_listenfd(char *port) {
  struct addrinfo hints, *listp, *p;
  int listenfd, rc, optval = 1;

  /* Get a list of potential server addresses */
  memset(&hints, 0, sizeof(struct addrinfo));
  hints.ai_socktype = SOCK_STREAM;             /* Accept connections */
  hints.ai_flags = AI_PASSIVE | AI_ADDRCONFIG; /* ... on any IP address */
  hints.ai_flags |= AI_NUMERICSERV;            /* ... using port number */
  if ((rc = getaddrinfo(NULL, port, &hints, &listp)) != 0) {
    fprintf(stderr, "getaddrinfo failed (port %s): %s\n", port,
            gai_strerror(rc));
    return -2;
  }

  /* Walk the list for one that we can bind to */
  for (p = listp; p; p = p->ai_next) {
    /* Create a socket descriptor */
    if ((listenfd = socket(p->ai_family, p->ai_socktype, p->ai_protocol)) < 0)
      continue; /* Socket failed, try the next */

    /* Eliminates "Address already in use" error from bind */
    setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, // line:netp:csapp:setsockopt
               (const void *)&optval, sizeof(int));

    /* Bind the descriptor to the address */
    if (bind(listenfd, p->ai_addr, p->ai_addrlen) == 0)
      break;                   /* Success */
    if (close(listenfd) < 0) { /* Bind failed, try the next */
      fprintf(stderr, "open_listenfd close failed: %s\n", strerror(errno));
      return -1;
    }
  }

  /* Clean up */
  freeaddrinfo(listp);
  if (!p) /* No address worked */
    return -1;

  /* Make it a listening socket ready to accept connection requests */
  if (listen(listenfd, LISTENQ) < 0) {
    close(listenfd);
    return -1;
  }
  return listenfd;
}
/* $end open_listenfd */

/****************************************************
 * Wrappers for reentrant protocol-independent helpers
 ****************************************************/

int Open_clientfd(char *hostname, char *port) {
  int rc;

  if ((rc = open_clientfd(hostname, port)) < 0)
    unix_error("Open_clientfd error");
  return rc;
}

int Open_listenfd(char *port) {
  int rc;

  if ((rc = open_listenfd(port)) < 0)
    unix_error("Open_listenfd error");
  return rc;
}

/*
 * Function prototypes
 */

void unix_error_w(char *msg) {
  fprintf(stderr, "%s: %s\n", msg, strerror(errno));
}

ssize_t Rio_readlineb_w(rio_t *rp, void *usrbuf, size_t maxlen) {
  ssize_t rc;

  if ((rc = rio_readlineb(rp, usrbuf, maxlen)) < 0) {
    unix_error_w("Rio_readlineb_w error");
    return 0;
  }
  return rc;
}

ssize_t Rio_readnb_w(rio_t *rp, void *usrbuf, size_t n) {
  ssize_t rc;

  if ((rc = rio_readnb(rp, usrbuf, n)) < 0)
    unix_error_w("Rio_readn_w error");
  return rc;
}

ssize_t Rio_writen_w(int fd, void *usrbuf, size_t n) {
  ssize_t written;

  if ((written = rio_writen(fd, usrbuf, n)) != n) {
    unix_error_w("Rio_writen_w error");
    return 0;
  }
  return written;
}

int Open_clientfd_w(char *hostname, char *port) {
  int rc;

  if ((rc = open_clientfd(hostname, port)) < 0)
    unix_error_w("Open_clientfd error");
  return rc;
}

ssize_t min(ssize_t a, ssize_t b) { return a > b ? b : a; }

ssize_t read_request_body(rio_t read_rio, char *buffer, size_t len) {
  // create the buffer
  char buf[MAXLINE + 1];

  // get the request/response header
  buf[0] = '\0';
  ssize_t content_length = 0, total_length = 0;
  while (strcmp(buf, "\r\n")) {
    ssize_t read_size = Rio_readlineb_w(&read_rio, buf, MAXLINE);
    if (!read_size) {
      return -1;
    }

    // accumulate size into a variable
    total_length += read_size;

    // check if the line contains Content-Length header
    if (strncmp(buf, "Content-Length: ", 16) == 0) {
      content_length = atoi(buf + 16);
    }
  }

  // if there's no request body
  if (content_length == 0) {
    return -1;
  }

  // read the request/responce body
  ssize_t current_pos = 0, read_size = 0;
  for (ssize_t i = 0; i < content_length; i += read_size) {
    // make the best use of the buffer
    read_size = Rio_readnb_w(
        &read_rio, buf + current_pos,
        min(min(MAXLINE - current_pos, content_length - i), MAX_READ_ONCE));
    if (read_size < 0) {
      return -1;
    } else if (!read_size) {
      break;
    }

    // increase current_pos, only if the buffer is full will the program write
    // the content into client
    current_pos += read_size;
    if (current_pos == MAXLINE) {
      break;
    }

    // accumulate size into a variable
    total_length += read_size;
  }
  // add the end of string
  buf[current_pos] = '\0';
  size_t length = min(current_pos, len - 1);

  // copy the content in buf to buffer
  memcpy(buffer, buf, length);
  buffer[length] = '\0';

  // return the bytes written
  return strlen(buffer);
}

void send_json_response(int client_fd, char *content, size_t len) {
  char buffer[MAXLINE];

  // Prepare the HTTP response headers
  const char *status_line = "HTTP/1.1 200 OK\r\n";
  const char *headers = "Content-Type: text/plain\r\n"
                        "Connection: close\r\n"
                        "Content-Length: ";

  // Calculate the content length
  int content_length = strlen(content);

  // Write the status line to buffer
  snprintf(buffer, MAXLINE, "%s", status_line);
  Rio_writen_w(client_fd, buffer, strlen(buffer));

  // Write the headers to buffer
  snprintf(buffer, MAXLINE, "%s%d\r\n\r\n", headers, content_length);
  Rio_writen_w(client_fd, buffer, strlen(buffer));

  // Write the JSON body to buffer
  snprintf(buffer, MAXLINE, "%s", content);
  Rio_writen_w(client_fd, buffer, strlen(buffer));

  // Flush the buffer to ensure all data is sent
  fsync(client_fd);
}

TEEC_Result send_message(size_t input_buf_len, const char *input_buffer,
                         size_t output_buf_len, char *output_buffer);

// function doit, nearly the same function as is in CSAPP
void doit(int listen_fd) {
  // some pre-defined variables
  char buf[MAXLINE], method[MAXLINE], uri[MAXLINE], version[MAXLINE];
  rio_t client_rio;

  int conn_fd = -1;
  socklen_t client_length = sizeof(struct sockaddr_storage);
  struct sockaddr_storage client_addr;

  // accept connection from client
  conn_fd = Accept(listen_fd, (SA *)&client_addr, &client_length);
  printf("Accepted connection from client\n");

  // read from client
  Rio_readinitb(&client_rio, conn_fd);
  if (!Rio_readlineb_w(&client_rio, buf, MAXLINE))
    goto end;
  if (sscanf(buf, "%s %s %s\r\n", method, uri, version) < 3 ||
      buf[strlen(buf) - 1] != '\n') {
    unix_error_w("Error occurred while parsing header line");
    goto end;
  }
  printf("Received request headline: %s", buf);

  // get the content sent from client
  char input_buf[MAXLINE] = {}, output_buf[MAXLINE] = {};
  if (read_request_body(client_rio, input_buf, MAXLINE) < 0) {
    unix_error_w("Error occurred while reading request body");
    goto end;
  }
  printf("Received request body: %s\n", input_buf);

  // get response string
  TEEC_Result result = send_message(MAXLINE, input_buf, MAXLINE, output_buf);
  if (result != TEEC_SUCCESS) {
    fprintf(stderr, "invoke failed, codes=0x%x\n", result);
  } else {
    printf("load TA successfully, TA's return: %s\n", output_buf);
  }

  // send the response to the client
  send_json_response(conn_fd, output_buf, MAXLINE);

end:
  // only close client
  if (conn_fd >= 0)
    Close(conn_fd);
}

static const TEEC_UUID g_demoTemplateUuid = {
    0x8e595120,
    0x7a1a,
    0x4d00,
    {0xa0, 0x8d, 0xc3, 0xa8, 0x48, 0xa4, 0xe0, 0x30}};

enum {
  CMD_GET_TA_VERSION = 1,
  CMD_BEGIN_SEND_MODEL_FILE = 2,
  CMD_SEND_MODEL_FILE_SEGMENT = 3,
  CMD_INFERENCE = 4,
  CMD_DELETE_MODEL_FILE = 5,
};

static TEEC_Context context = {0};
static TEEC_Session session = {0};
static TEEC_Operation operation = {0};
static uint32_t origin = 0;

TEEC_Result initialize(void) { return TEEC_InitializeContext(NULL, &context); }

TEEC_Result open_session(void) {
  /* MUST use TEEC_LOGIN_IDENTIFY method */
  operation.started = OPERATION_START_FLAG;
  operation.paramTypes =
      TEEC_PARAM_TYPES(TEEC_NONE, TEEC_NONE, TEEC_NONE, TEEC_NONE);

  return TEEC_OpenSession(&context, &session, &g_demoTemplateUuid,
                          TEEC_LOGIN_IDENTIFY, NULL, &operation, &origin);
}

TEEC_Result load_module_file(int file_size, const char *file_name) {
  TEEC_Result ret = TEEC_FAIL;

  operation.started = OPERATION_START_FLAG;
  operation.paramTypes =
      TEEC_PARAM_TYPES( // todo: what's the difference between
                        // TEEC_MEMREF_TEMP_INPUT and other variants?
          TEEC_VALUE_INPUT, TEEC_MEMREF_TEMP_INPUT, TEEC_NONE, TEEC_NONE);
  operation.params[0].value.a = file_size;
  operation.params[1].tmpref.size = strlen(file_name);
  operation.params[1].tmpref.buffer = (void *)file_name;

  ret = TEEC_InvokeCommand(&session, CMD_BEGIN_SEND_MODEL_FILE, &operation,
                           &origin);
  if (ret != TEEC_SUCCESS) {
    goto load_out1;
  }

  char file_buffer[8192];
  int file_fd = open(file_name, O_RDWR);
  if (file_fd < 0) {
    printf("%s: open %s failed!\n", __func__, file_name);
    ret = TEEC_FAIL;
    goto load_out1;
  }

  int copy_end = 0, read_ret = 0;
  while ((read_ret = read(file_fd, file_buffer, sizeof(file_buffer)))) {
    if (read_ret == -1) {
      break;
    }
    copy_end += read_ret;
    operation.started = OPERATION_START_FLAG;
    operation.paramTypes =
        TEEC_PARAM_TYPES( // todo: what's the difference between
                          // TEEC_MEMREF_TEMP_INPUT and other variants?
            TEEC_VALUE_INPUT, TEEC_MEMREF_TEMP_INPUT, TEEC_NONE, TEEC_NONE);
    operation.params[0].value.a = copy_end / sizeof(file_buffer);
    operation.params[1].tmpref.size = read_ret;
    operation.params[1].tmpref.buffer = file_buffer;
    ret = TEEC_InvokeCommand(&session, CMD_SEND_MODEL_FILE_SEGMENT, &operation,
                             &origin);
    if (ret != TEEC_SUCCESS) {
      goto load_out1;
    }
  }

load_out1:
  return ret;
}

TEEC_Result send_message(size_t input_buf_len, const char *input_buffer,
                         size_t output_buf_len, char *output_buffer) {
  operation.started = OPERATION_START_FLAG;
  operation.paramTypes = TEEC_PARAM_TYPES(
      TEEC_MEMREF_TEMP_INPUT, TEEC_MEMREF_TEMP_OUTPUT, TEEC_NONE, TEEC_NONE);
  operation.params[0].tmpref.buffer = (void *)input_buffer;
  operation.params[0].tmpref.size = input_buf_len;
  operation.params[1].tmpref.buffer = output_buffer;
  operation.params[1].tmpref.size = output_buf_len;
  return TEEC_InvokeCommand(&session, CMD_INFERENCE, &operation, &origin);
}

int main(int argc, char *argv[]) {
  TEEC_Result result;
  unsigned long file_size = -1;
  struct stat statbuff;

  /* Check arguments */
  if (argc != 3) {
    fprintf(stderr, "Usage: %s <port number> <file path>\n", argv[0]);
    exit(0);
  }

  // ignore SIGPIPE
  Signal(SIGPIPE, SIG_IGN);

  // open listen file descriptor
  int listen_fd = Open_listenfd(argv[1]);

  // check the status of the model file
  if (stat(argv[2], &statbuff) < 0) {
    printf("error! can't find model file\n");
    goto cleanup_1;
  } else {
    file_size = (int)statbuff.st_size;
    printf("module file size is %lu\n", file_size);
  }

  // initialize and open session
  result = initialize();
  if (result != TEEC_SUCCESS) {
    printf("Teec initialization failed!\n");
    goto cleanup_1;
  }

  result = open_session();
  if (result != TEEC_SUCCESS) {
    printf("teec open session failed!\n");
    goto cleanup_2;
  }

  // load the module file
  result = load_module_file(file_size, argv[2]);
  if (result != TEEC_SUCCESS) {
    printf("load_module_file failed, codes=0x%x\n", result);
    goto cleanup_3;
  } else {
    printf("load module file successfully\n");
  }

  // start listening and processing requests
  while (1) {
    doit(listen_fd);
  }

cleanup_3:
  TEEC_CloseSession(&session);
cleanup_2:
  TEEC_FinalizeContext(&context);
cleanup_1:
  return 0;
}
